import numpy as np
from copy import deepcopy
import cv2
import logging
from scipy.ndimage import gaussian_filter
from scipy.signal import savgol_filter


class Navpoint:
    """
    Navpoint
    基于图像的目标点生成
    """

    def __init__(self, min_area=7000, min_area_threshold=11000):
        """
        初始化路径规划器
        :param image: 图像分割结果，数字代表类别标签
        :param smoothness_threshold: 平滑度阈值，较高的值会选择平整的区域
        :param min_area: 选取的区域最小面积，避免选中噪声块
        :param min_area_threshold: 判断区域值为1的区域最小面积阈值
        """
        self.min_area = min_area
        self.min_area_threshold = min_area_threshold  # 新增参数

    def get_connected_components_and_centroids(self, image, category):
        """
        获取图像中指定类别的连通区域的标签和中心点坐标
        :param category: 类别值（1-5）
        :return: 每个连通区域的标签和中心点坐标
        """
        category_image = np.where(image == category, 1, 0).astype(np.uint8)
        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
            category_image, connectivity=4
        )  # 得到连通区域的数量、标签（0为背景，1-N为不同区域坐标）、连通区域坐标、连通区域质心坐标

        connected_components = []
        for label in range(1, num_labels):  # label=0 是背景，不处理
            # 计算区域内的点的数量（即该区域的面积）
            area = np.sum(labels == label)
            if area >= self.min_area:  # 过滤掉小区域
                center_x, center_y = centroids[label]
                connected_components.append(
                    (label, (int(center_x), int(center_y)), area)
                )

        return connected_components

    def calculate_smoothness(self, label):
        """
        计算区域的平滑程度，类别值越低，平滑度越高
        :param label: 区域标签
        :return: 平滑程度
        """
        smoothness = 6 - label  # 类别1 -> 平滑度高，类别5 -> 平滑度低
        return smoothness

    def check_for_obstacles(self, image, point1, point2):
        """
        检查区域值大于3的障碍区域是否存在于point1和point2之间
        参数：
            image  --- 分割图像
            point1 --- 第一个中心点坐标
            point2 --- 第二个中心点坐标
        返回:
            bool --- 是否存在障碍
        """
        # 获取这两个点之间的线段坐标
        x_coords = np.linspace(point1[0], point2[0], num=50).astype(int)
        y_coords = np.linspace(point1[1], point2[1], num=50).astype(int)

        # 确保坐标在界限内
        x_coords = np.clip(x_coords, 0, image.shape[1] - 1)
        y_coords = np.clip(y_coords, 0, image.shape[0] - 1)

        # 检测障碍
        for x, y in zip(x_coords, y_coords):
            if image[y, x] > 3 or image[y, x] == 0:  # 大于3或等于0则存在障碍
                return True

        return False

    def check_for_arrive(self, image, point):
        """
        检查目标点是否可到达
        参数：
            image --- 分割图像
            point --- 目标点坐标
        返回:
            bool  --- 目标点是否可到达
        """

        # 机器狗位置在图像底部中心 (x, y)
        height, width = image.shape
        bottom_center = (width // 2, height - 1)

        # 生成连接路径: 获取这两个点之间的线段坐标, 10 个坐标
        x_coords = np.linspace(bottom_center[0], point[0], num=10).astype(int)
        y_coords = np.linspace(bottom_center[1], point[1], num=10).astype(int)

        # 确保坐标在界限内
        x_coords = np.clip(x_coords, 0, image.shape[1] - 1)
        y_coords = np.clip(y_coords, 0, image.shape[0] - 1)

        # 检测是否可以到达
        # 切片操作：x_coords[:6] 和 y_coords[:6] 分别截取前 6 个元素（索引 0 到 5）
        for x, y in zip(x_coords[:6], y_coords[:6]):
            if image[y, x] > 3 or image[y, x] == 0:  # 大于3或等于0则存在障碍
                return False
        return True

    def midpoint(self, point1, point2):
        """
        计算两个点的中点坐标
        参数:
            point1 -- 第一个点的坐标，格式为(x1, y1)
            point2 -- 第二个点的坐标，格式为(x2, y2)
        返回:
            中点坐标，格式为(x, y)
        """
        x = int((point1[0] + point2[0]) / 2)
        y = int((point1[1] + point2[1]) / 2)
        return (x, y)

    def generate_horizontal_offsets(self, target_x, image_width=640, step=40):
        """
        生成水平偏移坐标, 以 target_x 为中心, 向左和向右偏移
        [500, 500 - 40, 500 + 40, 500 - 80, 500 + 80, ...]
        参数:
            target_x (int): 中心坐标
            image_width (int): 图像宽度
            step (int): 每次偏移的像素步长
        返回:
            list: 偏移坐标列表，从 target_x 开始，包含向左和向右的偏移
        """
        offsets = []
        max_offset = max(target_x, image_width - 1 - target_x)

        for i in range(step, max_offset + step, step):
            left = target_x - i
            right = target_x + i
            if 0 <= left < image_width:
                offsets.append(left)
            if 0 <= right < image_width:
                offsets.append(right)

        return [target_x] + offsets

    def select_target_region(self, res, point=None):
        """
        选择目标区域：根据面积和类别值选择最优区域
        :return: 选中的目标区域的中心点
        """
        if not isinstance(res, np.ndarray):
            seg_image = np.array(res.pred_sem_seg.data.squeeze().cpu())
        else:
            seg_image = res
        seg_image = seg_image.astype(np.uint8)
        assert np.max(seg_image) <= 5, f"存在类别{np.max(seg_image)}超过映射最大类别(5)"

        height, width = seg_image.shape

        if point is not None:
            # 机器狗位置在图像底部中心 (x, y)
            robot_point = (width // 2, height - 1)

            # 目标点是归一化的点，需要转换成实际坐标 (x, y)
            target_point = (int(point[0] * (width - 1)), int(point[1] * (height - 1)))
            logging.debug(f"robot_point: {robot_point}, target_point: {target_point}")

            offsets = self.generate_horizontal_offsets(target_point[0], width, step=40)

            # 检查偏移点是否存在障碍
            for x_offset in offsets:
                offset_point = (x_offset, target_point[1])
                if self.check_for_arrive(seg_image, offset_point) and seg_image[
                    offset_point[1], offset_point[0]
                ] in [1, 2, 3]:
                    return self.midpoint(offset_point, robot_point)

        best_region = None
        best_score = -float("inf")

        # 对每个类别（1/2/3）进行分析
        for category in range(1, 4):
            connected_components = self.get_connected_components_and_centroids(
                seg_image, category
            )

            for label, centroid, area in connected_components:
                # 处理区域值为1的区域
                if category == 1:
                    if area < self.min_area_threshold:  # 判断区域面积是否小于阈值
                        # 检查是否存在区域值为2的符合条件的区域
                        category2_components = (
                            self.get_connected_components_and_centroids(seg_image, 2)
                        )

                        has_no_obstacle = False

                        # 检查值为1的区域和值为2的区域之间是否存在障碍
                        centroid2_x = []
                        centroid2_y = []
                        for label2, centroid2, area2 in category2_components:
                            if not self.check_for_obstacles(
                                seg_image, centroid, centroid2
                            ):
                                has_no_obstacle = True
                                centroid2_x.append(centroid2[0])
                                centroid2_y.append(centroid2[1])

                        if has_no_obstacle:

                            # 如果没有值大于3的区域，则合并区域值为1和区域值为2
                            merged_centroid = (
                                int(centroid[0] + np.mean(centroid2_x)) // 2,
                                int(centroid[1] + np.mean(centroid2_y)) // 2,
                            )
                            centroid = merged_centroid  # 更新中点计算

                # 排除不可通行的区域（类别为4或5的区域）
                if category == 4 or category == 5:
                    continue

                smoothness = self.calculate_smoothness(category)

                # 加权计算区域的优选值，越大的值越好
                if category == 1:  # 平坦区域优先选择
                    score = area * smoothness * 2  # 平坦区域加权
                elif category == 2:  # 粗糙区域次之
                    score = area * smoothness  # 粗糙区域加权
                else:
                    score = area

                if score > best_score:
                    best_score = score
                    best_region = centroid

        if best_region is None:
            return None
        # 过滤点
        return best_region
